#!/usr/bin/python
# -*- encoding: utf8 -*-

import unittest
import random

import ds.heap as heap

class HeapTest(unittest.TestCase):
    def setUp(self):
        self.heap = heap.Heap()

    def testHeapEmpty(self):
        self.assertRaises(heap.HeapError, self.heap.pop)

    def testHeap(self):
        self.heap.push(10)
        self.assertEquals(1, self.heap.size())

        self.heap.push(20)
        self.assertEquals(2, self.heap.size())

        self.assertEquals(20, self.heap.pop())
        self.assertEquals(1, self.heap.size())
        self.assertEquals(10, self.heap.pop())
        self.assertEquals(0, self.heap.size())

    def testHeapWithOrderedData(self):
        n = 500
        data = range(n)
        data.reverse()
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i in data:
            self.assertEquals(i, self.heap.pop())

    def testHeapWithReversedData(self):
        n = 500
        data = range(n)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i in reversed(data):
            self.assertEquals(i, self.heap.pop())

    def testHeapWithShuffleData(self):
        n = 500
        data = range(n)
        random.shuffle(data)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i in sorted(data, reverse=True):
            self.assertEquals(i, self.heap.pop())

    def testHeapRemove(self):
        n = 500
        data = range(n)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i, obj in enumerate(data):
            self.heap.remove(obj)
            self.assertEquals(n - (i + 1), self.heap.size())

    def testHeapRemove2(self):
        n = 500
        data = range(n)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i, obj in enumerate(data[:100]):
            self.heap.remove(obj)
            self.assertEquals(n - (i + 1), self.heap.size())

        for i, obj in enumerate(data[400:]):
            self.heap.remove(obj)
            self.assertEquals(n - 100 - (i + 1), self.heap.size())

        data = data[100:400]
        self.assertEquals(len(data), self.heap.size())
        for i in reversed(data):
            self.assertEquals(i, self.heap.pop())

    def testHeapUpdate(self):
        n = 500
        data = range(n)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        new_data = range(1000, 1100)
        for obj, new_obj in zip(data[:100], new_data):
            self.heap.update(obj, new_obj)
            self.assertEquals(n, self.heap.size())

        new_data2 = range(-100, 0)
        for obj, new_obj in zip(data[400:], new_data2):
            self.heap.update(obj, new_obj)
            self.assertEquals(n, self.heap.size())

        data = data[100:400] + new_data + new_data2
        self.assertEquals(len(data), self.heap.size())
        for i in sorted(data, reverse=True):
            self.assertEquals(i, self.heap.pop())

class MinHeapTest(unittest.TestCase):
    def setUp(self):
        self.heap = heap.Heap(lambda x, y: cmp(y, x))

    def testHeapWithOrderedData(self):
        n = 500
        data = range(n)
        for i, obj in enumerate(data):
            self.heap.push(obj)
            self.assertEquals(i + 1, self.heap.size())

        for i in data:
            self.assertEquals(i, self.heap.pop())


if __name__ == '__main__':
    unittest.main()
